#启动训练过程
"""
参数解析和设备设置。
数据集的准备和解压。
实例化数据加载器、模型和实验构建器。
运行整个训练实验。
"""

import torch

from data import MetaLearningSystemDataLoader
from experiment_builder import ExperimentBuilder
from few_shot_learning_system import MAMLFewShotClassifier
from utils.parser_utils import get_args
from utils.dataset_tools import maybe_unzip_dataset
import  wandb  # 导入 wandb使结果可视化

#wandb.login(key="323911e8f5908bdecf4af73f5763ce113717bae8")

def main():

    # Combines the arguments, model, data and experiment builders to run an experiment
    args, device = get_args()
    # 指定 GPU ID，如果使用 Windows，可以直接在这里指定
    args.gpu_id = 0  # 假设使用第一个 GPU
    device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu")

    model = MAMLFewShotClassifier(args=args, device=device,
                                  im_shape=(2, 3,
                                            args.image_height, args.image_width))
    #maybe_unzip_dataset(args=args)
    data = MetaLearningSystemDataLoader(args=args, current_iter=0)
    maml_system = ExperimentBuilder(model=model, data=data, args=args, device=device)
    maml_system.run_experiment()

if __name__ == '__main__':
    main()